%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
%   Filename:   LCC_demons3
%   Author:     Jihun Oh
%   Date:       05-06-2014
%   Contact:    oj9040@gmail.com
%
% -------------------------------------------------------------------------
% This function performs image registration using Local Correlation Coefficient based Demons.   
% This can be applied to registering linearly correlated image volumes
% (e.g., contrast-enhanced image sequences)
%   Inputs:
%           S:      Static(Fixed) volume image
%           M:      Moving volume image
%           Mask:   Region where similarity is measured (0 or 1) (default=
%                   ones(size(S))
%           iter:   Maximum iteration (default=100)
%           diffeomorphism: diffeomorphic (default='off')
%   Outputs:
%           M2:     Deformed moving volume image           
%           Tx, Ty, Tz: Displacement fields in x, y, and z axis
%   
%   
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [M2, Tx, Ty, Tz] = LCC_demons3(S, M, Mask, voxelsize, iter, diffeomorphism)

%%%%%%%%%% Path %%%%%%%%%%
addpath('../Library');

if ~exist('Mask', 'var')
    Mask = ones(size(S));
end

if ~exist('voxelsize', 'var')
    voxelsize.x = 1;
    voxelsize.y = 1;
    voxelsize.z = 1;
end

if ~exist('iter', 'var')
    iter = 100;
end

if ~exist('diffeomorphism', 'var')
    diffeomorphism = 'off';
end

%%%%%%%% Parameters %%%%%%%%
% ALPHA (noise) constant
ALPHA = 0.5;   
% Gaussian Kernel for local statistics
SIG1 = 3;
% smoothing kernel for displacement incremental, Fluid-like regularization
SIG2 = 3;
% smoothing kernel for total displacement, Diffusion-like regularization
SIG3 = 3;

M2 = M;

% Initialize displacement field
Tx = zeros(size(S)); Ty = zeros(size(S)); Tz = zeros(size(S));
Txo = zeros(size(S)); Tyo = zeros(size(S)); Tzo = zeros(size(S));

[Sy,Sx,Sz] = gradient(S, voxelsize.y, voxelsize.x, voxelsize.z);
% S_g : Gaussian-Windowed convolution with S
S_g = anisoimgaussian3(S, SIG1, voxelsize.x, voxelsize.y, voxelsize.z);
% S2_g : Gaussian-Windowed convolution with S^2
S2_g = anisoimgaussian3(S.^2, SIG1, voxelsize.x, voxelsize.y, voxelsize.z);

ind = find(Mask ~= 0);

for itt=1:iter
    
    % Gradient of M
    [My, Mx, Mz] = gradient(M2, img2.y_voxel_size, img2.x_voxel_size, img2.z_voxel_size);
    % M_g : Gaussian-Windowed convolution with M
    M_g = anisoimgaussian3(M2, SIG1, img1.x_voxel_size, img1.y_voxel_size, img1.z_voxel_size);
    % SM_g : Gaussian-Windowed convolution with S*M
    SM_g = anisoimgaussian3(S.*M2, SIG1, img1.x_voxel_size, img1.y_voxel_size, img1.z_voxel_size);
    % M2_g : Gaussian-Windowed convolution with M^2
    M2_g = anisoimgaussian3(M2.^2, SIG1, img1.x_voxel_size, img1.y_voxel_size, img1.z_voxel_size);
    
    % The derivative of LCC with respect to displacment
    % Use the simplified LCC (sLCC)
    A1 = (S-S_g - (M2-M_g).*(SM_g-S_g.*M_g)./(M2_g-M_g.^2+eps))./(sqrt(S2_g-S_g.^2).*sqrt(M2_g-M_g.^2)+eps);
    dx1_sLCC = A1.*Mx;
    dy1_sLCC = A1.*My;
    dz1_sLCC = A1.*Mz;
    A2 = ((S-S_g).*(SM_g-S_g.*M_g)./(S2_g-S_g.^2+eps) - (M2-M_g))./(sqrt(M2_g-M_g.^2).*sqrt(S2_g-S_g.^2)+eps);
    dx2_sLCC = A2.*Sx;
    dy2_sLCC = A2.*Sy;
    dz2_sLCC = A2.*Sz;           
    clear A1 A2 Mx My Mz  M_g SM_g M2_g

    % Compute velocity fields in x, y and z coordinate
    B1 = 2*LCC./(dx1_sLCC.^2+dy1_sLCC.^2+dz1_sLCC.^2 + 4*ALPHA^2*LCC.^2 + eps);
    B2 = 2*LCC./(dx2_sLCC.^2+dy2_sLCC.^2+dz2_sLCC.^2 + 4*ALPHA^2*LCC.^2 + eps);
    Ux = B1.*dx1_sLCC + B2.*dx2_sLCC;
    Uy = B1.*dy1_sLCC + B2.*dy2_sLCC;
    Uz = B1.*dz1_sLCC + B2.*dz2_sLCC;
    clear B1 B2 dx1_SLCC dx2_SLCC dy1_SLCC dy2_SLCC dz1_SLCC dz2_SLCC CC

    % Smooth the velocity fields
    Ux = anisoimgaussian3(Ux, SIG2, voxelsize.x, voxelsize.y, voxelsize.z);
    Uy = anisoimgaussian3(Uy, SIG2, voxelsize.x, voxelsize.y, voxelsize.z);
    Uz = anisoimgaussian3(Uz, SIG2, voxelsize.x, voxelsize.y, voxelsize.z);

    % Add the velocity fields to the displacement fields
    Tx = Tx + Ux;
    Ty = Ty + Uy;
    Tz = Tz + Uz;  
    
    % Smooth the total displacement field
    Tx = anisoimgaussian3(Tx, SIG3, voxelsize.x, voxelsize.y, voxelsize.z);
    Ty = anisoimgaussian3(Ty, SIG3, voxelsize.x, voxelsize.y, voxelsize.z);
    Tz = anisoimgaussian3(Tz, SIG3, voxelsize.x, voxelsize.y, voxelsize.z);
    
    % Warping M
    M2 = warping3(M, X, Y, Z, X+Tx, Y+Ty, Z+Tz, 'linear');

    % Measure local correlation coefficients
    LCC = (SM_g - S_g.*M_g)./(sqrt(S2_g-S_g.^2).*sqrt(M2_g-M_g.^2)+eps);
    SLCC = sum(LCC(ind));
    
    % Measure displacement update
    diff = sqrt((Txo-Tx).^2 + (Tyo-Ty).^2 + (Tzo-Tz).^2);
    diff = sum(diff(ind))/numel(ind);
    
    display(sprintf('Iter = %d, SLCC = %3.5f, Diff. of Displacement = %3.5f', itt, SLCC, diff));
    
    % iteration condition
    if diff < 1e-3
        break;
    else
        Txo = Tx; Tyo = Ty; Tzo = Tz;
    end
        
end

